package stress

import (
	"context"
	"fmt"
	"gitee.com/jn-qq/simple-stress-testing/request"
	"gitee.com/jn-qq/simple-stress-testing/service"
	"gitee.com/jn-qq/simple-stress-testing/statistics"
	"log"
	"os"
	"os/signal"
	"runtime"
	"sync"
	"syscall"
	"time"
)

// Stress 压测结构体
type Stress struct {
	customerNum int           // 并发数
	requestNum  int           // 请求数(单个并发/协程)
	rSpaceTime  time.Duration // 每个请求间隔
	cpuNum      int           // CUP 核数，默认为一核，一般场景下单核已经够用了
	testTime    time.Duration // 压测程序最大执行时间(s)，默认不设置:0
	step        int           // 递增并发数
	stepTime    time.Duration // 递增并发数间隔时间
	requests    []*request.Request
}

// DefaultStress 默认压测对象
var DefaultStress = Stress{
	customerNum: 1,
	requestNum:  1,
	cpuNum:      runtime.NumCPU(),
	testTime:    0,
	step:        0,
	stepTime:    0,
	rSpaceTime:  time.Millisecond * 10,
}

// NewStress 创建压测对象
//
//	customerNum -> 初始开启线程数
//	requestNum  -> 每个线程请求次数，0 表示一直请求
//	cpuNum      -> 使用CPU核心数
//	step        -> 每次增加多少线程
//	stepTime    -> 每次增线程时间间隔
//	rSpaceTime  -> 每个线程中每次请求时间间隔
//	testTime    -> 压测总时间 0 表示一直压测
func NewStress(customerNum, requestNum, cpuNum, step int, stepTime, rSpaceTime, testTime time.Duration) *Stress {
	var stress Stress
	stress.copy(&DefaultStress)
	if customerNum > 1 {
		stress.customerNum = customerNum
	}
	if cpuNum > 1 {
		stress.cpuNum = cpuNum
	}
	if step > 0 && stepTime > 0 {
		stress.step = step
		stress.stepTime = stepTime
	} else if (step > 0 && stepTime == 0) || (step == 0 && stepTime > 0) {
		log.Fatal("并发数和并发间隔时间要同时指定")
		return nil
	}
	stress.testTime = testTime
	stress.requestNum = requestNum
	stress.rSpaceTime = rSpaceTime
	return &stress
}

func (s *Stress) Run(requests ...*request.Request) {

	if requests == nil || len(requests) == 0 {
		return
	} else {
		s.requests = requests
	}

	runtime.GOMAXPROCS(s.cpuNum)
	fmt.Printf("\n开始启动  并发数:%d 请求数:%d 请求参数: \n", s.customerNum, s.requestNum)
	for index, r := range requests {
		fmt.Printf("\t第%d步请求，如下：", index)
		r.Print()
	}

	// 开始处理
	// 上下文关系对象
	ctx := context.Background()
	// 判断执行时间
	if s.testTime > 0 {
		// 指定延时后关闭上下文
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, s.testTime)
		defer cancel()
		deadline, ok := ctx.Deadline()
		if ok {
			fmt.Printf(" deadline %s", deadline)
		}
	}

	// 处理 ctrl+c 信号
	ctx, cancelFunc := context.WithCancel(ctx)
	// 启动协程监听 ctrl+c 程序中断信号
	go func() {
		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGINT)
		<-c
		cancelFunc()
	}()

	var (
		wg  sync.WaitGroup // 发送数据完成
		wg1 sync.WaitGroup // 数据处理完成
	)

	// 收集者线程
	wg1.Add(1)
	statics := new(statistics.Statistics)
	statics.AddCustomerNum(s.customerNum)
	go statics.Receive(ctx, &wg1)

	// 请求线程
	for i := 0; i < s.customerNum; i++ {
		wg.Add(1)
		request.ConcurrencyMap.CNAdd(1)
		go service.Dispose(ctx, &wg, i, s.requestNum, s.rSpaceTime, requests...)
	}

	// 递增线程
	if s.step > 0 && s.stepTime > 0 {
		go s.stepCustomer(ctx, &wg, statics)
	}
	// 等待所有的数据都发送完成
	wg.Wait()
	// 延时1毫秒 确保数据都处理完成了
	time.Sleep(1 * time.Millisecond)
	close(request.CH)
	// 数据全部处理完成了
	wg1.Wait()
}

func (s *Stress) copy(m *Stress) {
	*s = *m
}

func (s *Stress) stepCustomer(ctx context.Context, wg *sync.WaitGroup, statics *statistics.Statistics) {
	// 设置定时
	ticker := time.NewTicker(s.stepTime)
	for {
		select {
		case <-ticker.C:
			// 添加协程数量
			for i := s.customerNum; i < s.customerNum+s.step; i++ {
				wg.Add(1)
				request.ConcurrencyMap.CNAdd(1)
				go service.Dispose(ctx, wg, i, s.requestNum, s.rSpaceTime, s.requests...)
				statics.AddCustomerNum(1)
			}
			s.customerNum += s.step
		case <-ctx.Done():
			// 处理完成
			return
		}
	}
}
